package org.vison.wonfu.framework.soa.discovery.LoadBalanceImpl;

import org.vison.wonfu.framework.soa.annotation.LoadBalance;
import org.vison.wonfu.framework.soa.discovery.LoadBalanceAlgorithm;
import org.vison.wonfu.framework.soa.discovery.ServiceInstance;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

public class ShortestResponseTimeAlgorithm implements LoadBalanceAlgorithm {

    private final ConcurrentMap<String, AtomicLong> totalResponseTimes = new ConcurrentHashMap<>();
    // 存储每个服务实例的请求次数 (实例ID -> 请求次数)
    private final ConcurrentMap<String, AtomicLong> requestCounts = new ConcurrentHashMap<>();

    private List<ServiceInstance> serviceInstances;

    public void setServiceInstances(List<ServiceInstance> serviceInstances) {
        this.serviceInstances = serviceInstances;
    }


    @Override
    public ServiceInstance select() {
        if (serviceInstances == null || serviceInstances.isEmpty()) {
            return null;
        }

        // 如果是第一次请求，随机选择一个实例
        if (totalResponseTimes.isEmpty()) {
            return serviceInstances.get(0);
        }

        ServiceInstance bestInstance = null;
        double minAverageTime = Double.MAX_VALUE;

        // 找出平均响应时间最短的实例
        for (ServiceInstance instance : serviceInstances) {
            String instanceId = instance.getServiceId();
            long totalTime = totalResponseTimes.getOrDefault(instanceId, new AtomicLong(0)).get();
            long count = requestCounts.getOrDefault(instanceId, new AtomicLong(0)).get();

            // 新实例或没有请求记录的实例优先考虑
            if (count == 0) {
                return instance;
            }

            double averageTime = (double) totalTime / count;

            // 找到平均响应时间最短的实例
            if (averageTime < minAverageTime) {
                minAverageTime = averageTime;
                bestInstance = instance;
            }
        }

        return bestInstance != null ? bestInstance : serviceInstances.get(0);
    }

    /**
     * 记录请求响应时间，用于后续计算平均值
     */
    public void recordResponseTime(ServiceInstance instance, long responseTimeMs) {
        String instanceId = instance.getServiceId();

        // 累加总响应时间
        totalResponseTimes.computeIfAbsent(instanceId, k -> new AtomicLong(0))
                .addAndGet(responseTimeMs);

        // 累加请求次数
        requestCounts.computeIfAbsent(instanceId, k -> new AtomicLong(0))
                .incrementAndGet();
    }

    @Override
    public LoadBalance getType() {
        return LoadBalance.FASTEST_RESPONSE;
    }
}
